{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 安装的前提就是你的系统需要是linux系统并且按照步骤已经装好了CUDA和CUDNN\n",
    "git clone https://github.com/NVIDIA/apex\n",
    "cd apex\n",
    "pip install -v --no-cache-dir --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" ./"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'xxxNet' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "\u001B[0;32m<ipython-input-2-98edb3a82fef>\u001B[0m in \u001B[0;36m<module>\u001B[0;34m\u001B[0m\n\u001B[1;32m      2\u001B[0m \u001B[0;32mimport\u001B[0m \u001B[0mtorch\u001B[0m\u001B[0;34m,\u001B[0m\u001B[0mtorchvision\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m      3\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m----> 4\u001B[0;31m \u001B[0mnet\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mxxxNet\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m      5\u001B[0m \u001B[0mnet\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mcuda\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m      6\u001B[0m \u001B[0mnet\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mtrain\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;31mNameError\u001B[0m: name 'xxxNet' is not defined"
     ]
    }
   ],
   "source": [
    "import apex.amp as amp\n",
    "import torch,torchvision\n",
    "\n",
    "net = xxxNet()\n",
    "net.cuda()\n",
    "net.train()\n",
    "\n",
    "params_low_lr = []\n",
    "params_high_lr = []\n",
    "for n, p in net.named_parameters():\n",
    "    if 'encoder' in n:\n",
    "        params_low_lr.append(p)\n",
    "    else:\n",
    "        params_high_lr.append(p)\n",
    "opt = Adam([{'params': params_low_lr, 'lr': 5e-5},\n",
    "                   {'params': params_high_lr, 'lr': 1e-4}], weight_decay=settings.WEIGHT_DECAY)\n",
    "net, opt = amp.initialize(net, opt, opt_level=\"O1\") # 这里要添加这句代码\n",
    "...\n",
    "...\n",
    "loss = ...\n",
    "if not torch.isnan(loss):\n",
    "    opt.zero_grad()\n",
    "    with amp.scale_loss(loss, opt) as scaled_loss:\n",
    "        scaled_loss.backward()   # loss要这么用\n",
    "    opt.step()\n",
    "\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 如果需要使用resume方式训练网络的时候就可以使用如下的代码"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'model' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "\u001B[0;32m<ipython-input-4-c82474a43d42>\u001B[0m in \u001B[0;36m<module>\u001B[0;34m\u001B[0m\n\u001B[1;32m      1\u001B[0m \u001B[0;31m# Initialization\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m      2\u001B[0m \u001B[0mopt_level\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0;34m'O1'\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m----> 3\u001B[0;31m \u001B[0mmodel\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0moptimizer\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mamp\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0minitialize\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mmodel\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0moptimizer\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0mopt_level\u001B[0m\u001B[0;34m=\u001B[0m\u001B[0mopt_level\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m      4\u001B[0m \u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m      5\u001B[0m \u001B[0;31m# Train your model\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;31mNameError\u001B[0m: name 'model' is not defined"
     ]
    }
   ],
   "source": [
    "# Initialization\n",
    "opt_level = 'O1'\n",
    "model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)\n",
    "\n",
    "# Train your model\n",
    "...\n",
    "\n",
    "# Save checkpoint\n",
    "checkpoint = {\n",
    "    'model': model.state_dict(),\n",
    "    'optimizer': optimizer.state_dict(),\n",
    "    'amp': amp.state_dict()\n",
    "}\n",
    "torch.save(checkpoint, 'amp_checkpoint.pt')\n",
    "...\n",
    "\n",
    "# Restore\n",
    "model = ...\n",
    "optimizer = ...\n",
    "checkpoint = torch.load('amp_checkpoint.pt')\n",
    "\n",
    "model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)\n",
    "model.load_state_dict(checkpoint['model'])  # 注意，load模型需要在amp.initialize之后！！！\n",
    "optimizer.load_state_dict(checkpoint['optimizer'])\n",
    "amp.load_state_dict(checkpoint['amp'])\n",
    "\n",
    "# Continue training\n",
    "..."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "invalid syntax (<ipython-input-5-4a6fddf2962a>, line 1)",
     "output_type": "error",
     "traceback": [
      "\u001B[0;36m  File \u001B[0;32m\"<ipython-input-5-4a6fddf2962a>\"\u001B[0;36m, line \u001B[0;32m1\u001B[0m\n\u001B[0;31m    amp.initialize(models=,optimizers=,opt_level=)\u001B[0m\n\u001B[0m                          ^\u001B[0m\n\u001B[0;31mSyntaxError\u001B[0m\u001B[0;31m:\u001B[0m invalid syntax\n"
     ]
    }
   ],
   "source": [
    "amp.initialize(models=,optimizers=,opt_level=)\n",
    "# 其中opt_level有四种模式：\n",
    "\"\"\"\n",
    "* 'O0':纯FP32训练，可以作为accuracy的baseline\n",
    "* 'O1':混合精度训练，推荐！根据黑白名单自动决定使用FP16(GEMM，卷积)还是FP32训练(Softmax)进行训练\n",
    "* 'O2':几乎FP16训练\n",
    "* ‘03’:纯FP16训练，很不稳定，但是可以用作speed的baseline\n",
    "\"\"\""
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# 可以使用以下的函数来显示决定是否应该使用半精度浮点数来进行运算。\n",
    "amp.register_float_function(module=torch,name='softmax')\n",
    "amp.register_half_function(mudule=torch,name='conv')\n",
    "amp.register_promote_function(module,name=)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": 6,
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "invalid syntax (<ipython-input-6-94a1b42f2a91>, line 4)",
     "output_type": "error",
     "traceback": [
      "\u001B[0;36m  File \u001B[0;32m\"<ipython-input-6-94a1b42f2a91>\"\u001B[0;36m, line \u001B[0;32m4\u001B[0m\n\u001B[0;31m    amp.register_promote_function(module,name=)\u001B[0m\n\u001B[0m                                              ^\u001B[0m\n\u001B[0;31mSyntaxError\u001B[0m\u001B[0;31m:\u001B[0m invalid syntax\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 多GPU，同步BN，自动混合精度\n",
    "#### 一般情况下的训练代码\n",
    "* 导入数据结构到dataloader上\n",
    "* 定义模型迁移到gpu上\n",
    "* 定义损失函数和优化器\n",
    "* 把模型用nn.dataParallel迁移到gpu上"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import argparse\n",
    "import torch\n",
    "from apex import amp\n",
    "from apex.parallel import convert_syncbn_model\n",
    "from apex.parallel import DistributedDataParallel as DDP\n",
    "# 罗列以下自己网络的超参数\n",
    "def parse():\n",
    "    parser = argparse.ArgumentParser()\n",
    "    # local_rand很重要，其指定来输出的设备\n",
    "    parser.add_argument('--local_rank', type=int, default=0)\n",
    "    ...\n",
    "    ...\n",
    "    args = parser.parse_args()\n",
    "    return args\n",
    "# 在主函数的开头可以这样写\n",
    "def main():\n",
    "    args = parse()\n",
    "    torch.cuda.set_device(args.local_rank)\n",
    "    # torch.utils.launch也需要set_device所以必须写\n",
    "    import torch.distributed as dist\n",
    "    dist.init_process_group(backend='nccl',init_method='env://')\n",
    "args = parse(0)\n",
    "# 导入数据的借口有点不一样，需要使用一个DistributedSampler\n",
    "dataset = torchvision.datasets.MNIST('./MNIST',train=True)\n",
    "num_workers = 4 if cuda else 0\n",
    "# 多了一个DistributedSampler,作为dataloader的sampler\n",
    "train_sampler = torch.utils.data.distributed.DistributedSampler()\n",
    "loader = torch.utils.data.DataLoader(dataset,batch_size=args.batchsize\n",
    "                                     ,shuffle=False\n",
    "                                     ,num_workers=args.num_workers\n",
    "                                     ,pin_memory=cuda\n",
    "                                     ,drop_last=True,sampler=train_sampler)\n",
    "# 之后定义模型：\n",
    "net = XXXNet(using_amp=True)\n",
    "net.train()\n",
    "net = convert_syncbn_model(net)     # 使用apex支持的方法，使得普通的bn成为同步bn\n",
    "# 切记在网络实现中，不要使用torch中自带的SyncBatchnorm\n",
    "device = torch.device('cuda:{}'.format(args.local_rank))\n",
    "\n",
    "# 定义优化器，损失函数，定义优化器一定要在把模型搬运到GPU之后\n",
    "opt = torch.optim.Adam(params=[{'params':params_low_lr,'lr':4e-5},{'params':params_high_lr,'lr':1e-4}]\n",
    "                       ,weight_decay=0)\n",
    "crit = torch.nn.BCELoss().to(device)\n",
    "\n",
    "# 多GPU设置\n",
    "net,opt = amp.initalize(net,opt,opt_level='O1')\n",
    "net = DDP(net,delay_allreduce=True)\n",
    "\n",
    "# 记得loss要这样子使用:\n",
    "opt.zero_grad()\n",
    "# loss.backward()\n",
    "with amp.scale_loss(loss,opt) as scaled_loss:\n",
    "    scaled_loss.backward()\n",
    "opt.step()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% 首先先导入自己需要的包\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "这样子基本就完成了，但是不能在主函数中执行，无论是apex支持的DDP还是pytorch自身支持的DDP(torch.nn.parallel.DistributedDataParallel))\n",
    "都需要使用torch.distributed.launch来使用方法是：\n",
    "CUDA_VISIBLE_DEVICES=1,2,4 python -m torch.distributed.launch --nproc_per_node=3 train.py\n",
    "\n",
    "# 注意1,2,3是想用的GPU编号，nproc_per_node 是指定你用了几块GPU。nproc是开启几个进程，设置为\n",
    "# 和GPU一样的大小，意思是每一个进程需要负责一个GPU，per_Node代表了你只有一个主机的服务器\n",
    "#开头必须要加入loca_rank 是因为torch.distributed.launch会调用这个local_rank\n",
    "\n",
    "# resume模型需要设置在to(device)之前，就是以CPU导入模型才可以，否则可能会报一个显存不够的错误\n",
    "#于是resume可以用下面的简单代码来表示"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch\n",
    "# 定义模型\n",
    "net = torchvision.models.resnet18(pretrained=True,using_amp=True)\n",
    "net.train()\n",
    "# resume 旧的训练模型\n",
    "resume = './model_ranet_17/65000.pth'\n",
    "checkpoint = torch.load(resume,map_location='cpu')\n",
    "net.load_state_dict((checkpoint['model']))\n",
    "# 转同步bn\n",
    "net = convert_syncbn_model(net)\n",
    "\n",
    "# 搬到GPU上面取，local_rank虽然默认是0，但是torch.distributed.launch会为每一个进程分配单独的GPU\n",
    "# local_rank就是被这个lauch脚本重新设定了\n",
    "devcie = torch.device('cuda:{}'.format(args.local_rank))\n",
    "net = net.to(device)\n",
    "# 定义优化器\n",
    "params_low_lr = []\n",
    "params_high_lr = []\n",
    "for n,p in net.named_parameters():\n",
    "    if 'encoder' in n:\n",
    "        params_low_lr.append(p)\n",
    "    else:\n",
    "        params_high_lr.append(p)\n",
    "opt = torch.optim.Adam(params=[{'params':params_low_lr,'lr':4e-5},{'params':params_high_lr,'lr':1e-4}]\n",
    "                       ,weight_decay=0)\n",
    "crit = torch.nn.BCELoss().to(device)\n",
    "\n",
    "# 分布式训练：\n",
    "net,opt = amp.initialize(net,opt,opt_level='O0')\n",
    "net = DDP(net,delay_allreduce=True)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "保存模型需要注意的地方\n",
    "分布式训练开启多个进程，如果在代码中写了torch.save来保存模型，那么每一个进程都会保存一个限购听的模型\n",
    "就很容易导致模型的无法使用。。\n",
    "可以用如下的代码来保存模型，只用loca_rank==0来仅仅在第一个GPU上执行的进程来保存模型，就不会有文件重新读写的情况发生了"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "if args.locak_rank == 0:\n",
    "    torch.save(XXXX)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}